#!/usr/bin/env python3
# Author: Armit
# Create Time: 2022/11/18 

from argparse import ArgumentParser

import numpy as np
from sklearn.model_selection import train_test_split
from sklearn.neighbors import KNeighborsClassifier

from data import get_data, FEATURE_CAT, FEATURE_NUM
from utils import show_clf_metrics


def dist_fn(a:np.ndarray, b:np.ndarray) -> float:
  return np.square(a - b).sum()


def knn(args):
  X, y = get_data(args.limit, FEATURE_NUM)
  
  X_train, X_test, y_train, y_test = train_test_split(X, y, stratify=y, test_size=0.7, random_state=42)
  print(f'dataset: {len(X_train)} for train. {len(X_test)} samples for test')
  
  # 构建一个字典储存Undefined所在的行，这样就可以用.drop去掉它们
  Undefined = [i for i in y_train.index if y_train.loc[i] == 3]
  X_train = X_train.drop(labels=Undefined)
  y_train = y_train.drop(labels=Undefined)
  # breakpoint()
  print(f'[kNN] k={args.k} p={args.p}')
  knn = KNeighborsClassifier(n_neighbors=args.k, p=args.p, n_jobs=4)

  knn.fit(X_train, y_train)
  y_pred = knn.predict(X_test) # 试图用3分类训练的模型预测Undefined，虽然说可以得到结果

  # show_clf_metrics(y_test, y_pred)
  breakpoint()

if __name__ == '__main__':
  parser = ArgumentParser()
  parser.add_argument('-k', default=5, type=int, help='kNN n_neighbors')
  parser.add_argument('-p', default=2, type=int, help='minkowski distance power')
  # 因为类型增加，所以采样点数提高到40000
  parser.add_argument('-N', '--limit', default=40000, type=int, help='limit dataset size')
  args = parser.parse_args()

  knn(args)
